#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python常用模块库 
@File ：memory_block.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/2/20 上午12:48 
"""

import mmap
import os
from typing import Optional


class MemoryBlock:
	"""底层共享内存块的实现。

	Args:
		name (str): 共享内存的文件路径。
		size (int): 内存块大小（字节）。
		create (bool): 是否创建新内存块。
	"""

	def __init__(self, name: str, size: int, create: bool = True) -> None:
		self.name = name
		self.size = size
		self.create = create
		self.fd = None
		self.mm = None
		self._open()

	def _open(self) -> None:
		"""打开或创建共享内存文件。"""
		flags = os.O_CREAT | os.O_RDWR if self.create else os.O_RDWR
		self.fd = os.open(self.name, flags)
		if self.create:
			os.ftruncate(self.fd, self.size)
		self.mm = mmap.mmap(self.fd, self.size, access=mmap.ACCESS_WRITE)

	def write(self, data: bytes, offset: int = 0) -> None:
		"""写入数据到内存块。

		Args:
			data (bytes): 要写入的字节数据。
			offset (int): 写入偏移量。
		"""
		self.mm.seek(offset)
		self.mm.write(data)

	def read(self, offset: int = 0, size: Optional[int] = None) -> bytes:
		"""从内存块读取数据。

		Args:
			offset (int): 读取偏移量。
			size (int, optional): 读取字节数。

		Returns:
			bytes: 读取的字节数据。
		"""
		self.mm.seek(offset)
		if size is None:
			size = self.size - offset
		return self.mm.read(size)

	def close(self) -> None:
		"""关闭内存块资源。"""
		if self.mm:
			self.mm.close()
		if self.fd is not None:
			os.close(self.fd)
			self.fd = None

	def unlink(self) -> None:
		"""删除共享内存文件。"""
		if os.path.exists(self.name):
			os.unlink(self.name)
